from typing import Dict, List

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
from loguru import logger

router = APIRouter(prefix="/ws", tags=["websocket"])


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, List[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, username: str):
        await websocket.accept()
        if username not in self.active_connections:
            self.active_connections[username] = []
        if websocket not in self.active_connections[username]:
            self.active_connections[username].append(websocket)
        logger.info(f"User {username} connected. Total: {len(self.active_connections[username])}")

    def disconnect(self, websocket: WebSocket, username: str):
        if username in self.active_connections:
            try:
                self.active_connections[username].remove(websocket)
                logger.info(f"User {username} disconnected. Remaining: {len(self.active_connections[username])}")
            except ValueError:
                logger.warning(f"WebSocket not found for user {username} during disconnect.")
            if not self.active_connections[username]:
                del self.active_connections[username]
                logger.info(f"All connections for user {username} closed.")

    async def broadcast(self, message: str):
        to_remove = []
        for username, connections in self.active_connections.items():
            for connection in connections:
                try:
                    await connection.send_text(message)
                except Exception as e:
                    logger.warning(f"Broadcast failed for user {username}: {e}")
                    to_remove.append((username, connection))
        for username, connection in to_remove:
            self.disconnect(connection, username)

    async def send_to_users(self, usernames: List[str], message: str):
        for username in usernames:
            connections = self.active_connections.get(username, [])
            to_remove = []
            for connection in connections:
                try:
                    await connection.send_text(message)
                except Exception as e:
                    logger.warning(f"Send failed for user {username}: {e}")
                    to_remove.append(connection)
            for connection in to_remove:
                self.disconnect(connection, username)


manager = ConnectionManager()


@router.websocket("/broadcast")
async def websocket_endpoint(websocket: WebSocket, username: str = Query(...)):
    await manager.connect(websocket, username)
    try:
        while True:
            await websocket.receive_text()  # 保持连接，可忽略内容
    except WebSocketDisconnect:
        manager.disconnect(websocket, username)
        logger.info(f"WebSocketDisconnect for user {username}")
    except Exception as e:
        manager.disconnect(websocket, username)
        logger.error(f"Unexpected error for user {username}: {e}")
